import csv
import sqlite3
from typing import List

import requests
from langchain_chroma import Chroma
from langchain_community.utilities import SQLDatabase
from langchain_core.tools import tool
from pydantic import Field

from utils import BGE_Embed


@tool("get_weather")  # 工具的名称（注意这是必写内容）
def get_weather(location: str) -> str:
    """ 返回对应城市的天气信息 """  # 工具的描述（注意这是必写内容）
    print("查询天气信息中....")
    if location == "北京" or location == "北京市" or location.lower() == "beijing":
        return "北京 多云转阴 33/25℃"
    if location == "重庆" or location == "重庆市" or location.lower() == "chongqing":
        return "重庆 阴天 33/25℃"
    if location == "上海" or location == "上海市" or location.lower() == "shanghai":
        return "上海 晴 33/25℃"
    if location == "天津" or location == "天津市" or location.lower() == "tianjin":
        return "天津 小雨 33/25℃"


@tool("get_goods")
def get_goods(name: str) -> str:
    """ 查询商品信息 """
    print("查询商品中....")
    message = ""
    conn = sqlite3.connect("factory.db")
    cursor = conn.cursor()
    cursor = cursor.execute('SELECT id, name, price, number  from GOODS WHERE name=?', [name])
    for row in cursor:
        message += f'"ID": {row[0]}, "NAME": {row[1]}, "PRICE": {row[2]}, "NUMBER": {row[3]} \n'

    if message == "":
        message = "你所指定的商品，未查到"
    conn.commit()
    conn.close()

    return message


import xlwt


@tool("save_goods")
def save_goods(id: int, name: str, price: float, number: int):
    """ 保存查询的商品信息到excel中 """
    print("保存商品信息中....")
    workbook = xlwt.Workbook(encoding="utf-8")
    worksheet = workbook.add_sheet("商品数据")
    worksheet.write(0, 0, "商品编号")
    worksheet.write(0, 1, "商品名字")
    worksheet.write(0, 2, "商品价格")
    worksheet.write(0, 3, "商品数量")
    worksheet.write(1, 0, id)
    worksheet.write(1, 1, name)
    worksheet.write(1, 2, price)
    worksheet.write(1, 3, number)
    workbook.save("goods.xls")

    return "商品信息已保存至./goods.xls"


@tool("rag_search")
def rag_search(query: str):
    """ 检索三国演义书籍的上下文 """
    print("检索《三国演义》中 ...")
    embedding = BGE_Embed()
    vector_db = Chroma(
        embedding_function=embedding,
        persist_directory='../LangChain/data/chroma_db'
    )
    context = vector_db.similarity_search_with_score(query)  # 针对内容进行检索
    return context[0][0].__dict__['page-content']


if __name__ == '__main__':
    print(get_weather("重庆"))
